#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import math
# from diff_gauss import GaussianRasterizationSettings, GaussianRasterizer
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from utils.sh_utils import eval_sh
import numpy as np


def render(viewpoint_camera, 
           pc : GaussianModel, 
           pipe, 
           bg_color : torch.Tensor, 
           scaling_modifier=1.0, 
           override_color=None,
           transfer=True):
    
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:    screenspace_points.retain_grad()
    except: pass

    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
    image_height = int(viewpoint_camera.image_height)
    image_width = int(viewpoint_camera.image_width)

    raster_settings = GaussianRasterizationSettings(
        image_height=image_height,
        image_width=image_width,
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug)

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity
    scales = pc.get_scaling
    rotations = pc.get_rotation
    
    visible_mask = rasterizer.markVisible(positions=means3D)
    colors = pc.feature_encoder(viewpoint_camera, visible_mask)
    
    if transfer:
        colors_offsets, opacity_offsets, scales_offsets, rotations_offsets = pc.offsets_feature_encoder(viewpoint_camera, visible_mask)
        colors = colors + colors_offsets
        opacity = opacity[visible_mask] + opacity_offsets
        scales = scales[visible_mask] * scales_offsets
        rotations = pc.rotation_activation(rotations[visible_mask] + rotations_offsets)
        
        # data = {
        #     "color": colors.detach().cpu().numpy(),
        #     "opacity": opacity.detach().cpu().numpy(),
        #     "scale": scales.detach().cpu().numpy(),
        #     "rots": rotations.detach().cpu().numpy(),
        # }
        # array_data = np.array(list(data.items()))
        # np.save(f'ratio-{str(viewpoint_camera.ratio)[0]}.npy', array_data)
        
    else:
        # colors = pc.feature_encoder(viewpoint_camera, visible_mask)
        opacity = opacity[visible_mask]
        scales = scales[visible_mask]
        rotations = rotations[visible_mask]

    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    # rendered_image, radii, point_id, point_weight = rasterizer(means3D=means3D[visible_mask],
    rendered_image, radii = rasterizer(means3D=means3D[visible_mask],
                                        means2D=means2D[visible_mask],
                                        shs=None, 
                                        colors_precomp=colors,
                                        opacities=opacity,
                                        scales=scales,
                                        rotations=rotations,
                                        cov3D_precomp=None)
                            
    tem_radii = torch.zeros_like(visible_mask, dtype=radii.dtype, device=visible_mask.device)
    tem_radii[visible_mask] = radii
    visibility_filter = torch.logical_and(visible_mask, tem_radii>0)
        
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "visibility_filter" : visibility_filter,
            # "scaling": scales[visible_mask],
            "radii": tem_radii}
